package com.crazyteam.datasource.dynamic;

import org.apache.commons.lang.StringUtils;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.beans.factory.BeanFactoryAware;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.util.CollectionUtils;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author guoyang
 * @Date 2018-01-09
 * @Desc 动态数据源配置bean工厂
 */
public class DynamicDataSource extends AbstractRoutingDataSource implements BeanFactoryAware {

    private BeanFactory beanFactory;

    private DataSource masterDataSource; //主库=

    private List<DataSource> slavesDataSources; //备库组

    // 可用数据源
    private final Map<String, DataSource> supportedDataSources = new ConcurrentHashMap<>();
    private final Map<String, DataSource> unSupportedDataSources = new ConcurrentHashMap<>();

    private final byte[] switchLock = new byte[0];

    private boolean smartSwitch = false;

    private Integer checkDelay = 60;

    private String masterName;
    private List<String> slaveNames;

    DynamicDataSource(String masterName, List<String> slaveNames){
        this.masterName = masterName;
        this.slaveNames = slaveNames;
    }

    public void setCheckDelay(Integer checkDelay) {
        if (checkDelay != null){
            this.checkDelay = checkDelay;
        }
    }

    @Override
    public void afterPropertiesSet() {
        if (this.masterName == null) {
            throw new IllegalArgumentException("Property 'masterDataSource' is required");
        }
        masterDataSource = beanFactory.getBean(masterName, DataSource.class);
        if (this.masterName == null) {
            throw new IllegalArgumentException("Property 'masterDataSource' is required");
        }
        if (!CollectionUtils.isEmpty(slaveNames)){
            slavesDataSources = new ArrayList<>();
            slaveNames.forEach((slave) -> slavesDataSources.add(Optional.of(beanFactory.getBean(slave, DataSource.class)).get()));
        }
        setDefaultTargetDataSource(masterDataSource);
        Map<Object, Object> targetDataSources = new HashMap<>();
        targetDataSources.put(DynamicDataSourceHolder.MASTER, masterDataSource);
        supportedDataSources.put(DynamicDataSourceHolder.MASTER, masterDataSource);
        for(int idx = 0 ; idx < slavesDataSources.size(); idx++){
            targetDataSources.put(DynamicDataSourceHolder.SLAVE + idx, slavesDataSources.get(idx));
            supportedDataSources.put(DynamicDataSourceHolder.SLAVE + idx, slavesDataSources.get(idx));
        }
        setTargetDataSources(targetDataSources);
        super.afterPropertiesSet();

        // 如果有数据库动态移出，则需要数据库定时校验还原
        if (smartSwitch){
            new Timer().schedule(new TimerTask() {
                @Override
                public void run() {
                    unSupportedDataSources.forEach((key, ds) -> {
                        try {
                            Connection conn = ds.getConnection();
                            if (conn != null){
                                synchronized (switchLock){
                                    unSupportedDataSources.remove(key);
                                    supportedDataSources.put(key, ds);
                                }
                                logger.warn("数据库" + key + "恢复可用");
                                try{
                                    conn.close();
                                } catch (SQLException e) {
                                    logger.error("", e);
                                }
                            }else {
                                logger.warn("数据库" + key + "依旧不可用");
                            }
                        } catch (SQLException e) {
                            logger.warn("数据库" + key + "依旧不可用");
                            // do nothing
                        }
                    });
                }
            }, checkDelay * 1000L, 10000L);
        }
    }

    @Override
    protected String determineCurrentLookupKey() {

        String dataSourceName = DynamicDataSourceHolder.getDataSource();

        String supportReadWriter = System.getProperty("DynamicDataSource.supportReadWriter");
        if (StringUtils.isNotEmpty(supportReadWriter) && supportReadWriter.equalsIgnoreCase("false")){
            return DynamicDataSourceHolder.MASTER;
        }
        String sourceName = supportedDataSources.keySet().stream().filter((key) -> key.startsWith(dataSourceName)).reduce((k1, k2) -> new Random().nextBoolean()? k1:k2).orElse(DynamicDataSourceHolder.MASTER);
        if(StringUtils.isNotEmpty(sourceName) && supportedDataSources.containsKey(sourceName)) {
            return sourceName;
        }else {
            return DynamicDataSourceHolder.MASTER;
        }
    }

    /**
     * 获取数据源，如果获取失败 则切换数据源重新获取
     * @param lookupKey
     * @return
     * @throws SQLException
     */
    private Connection getConnection(String lookupKey) throws SQLException {
        DataSource dataSource = this.supportedDataSources.get(lookupKey);
        if (dataSource == null || DynamicDataSourceHolder.MASTER.equalsIgnoreCase(lookupKey)) {
            return masterDataSource.getConnection();
        } else {
            try {
                Connection conn = dataSource.getConnection();
                if (conn == null) {
                    return switchConnection(lookupKey, dataSource);
                }else {
                    return conn;
                }
            } catch (Exception e) {
                logger.error("数据库" + lookupKey + "获取链接失败,切换备库", e);
                return switchConnection(lookupKey, dataSource);
            }
        }
    }

    /**
     * 切换 数据源，并移出不可用数据源
     * @param oldKey
     * @param oldDataSource
     * @return
     * @throws SQLException
     */
    private Connection switchConnection(String oldKey, DataSource oldDataSource) throws SQLException{
        // 将不支持的数据源移除，重新设置数据源
        synchronized (switchLock) {
            if (supportedDataSources.containsKey(oldKey)){
                this.supportedDataSources.remove(oldKey);
            }
            if (!unSupportedDataSources.containsKey(oldKey)){
                this.unSupportedDataSources.put(oldKey, oldDataSource);
            }
        }
        String newLookupKey = determineCurrentLookupKey();
        logger.warn("成功切换数据源：" +  oldKey + "切换到：" + newLookupKey);
        return getConnection(newLookupKey);
    }

    /**
     * 重写获取连接接口，如果要动态切换，则配置 smartSwitch=true
     * 配置后。或校验获取的 连接报错或空连接，则将数据源切换到不可用
     * 用异步线程去校验，当可正常获取连接是，在开启
     * 配置需要配置数据库连接属性 testOnBorrow=true 才有效
     * @return
     * @throws SQLException
     */
    @Override
    public Connection getConnection() throws SQLException {
        if (smartSwitch){
            String lookupKey = determineCurrentLookupKey();
            return getConnection(lookupKey);
        }else {
            return super.getConnection();
        }
    }

    @Override
    public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
        this.beanFactory = beanFactory;
    }

    public void setSmartSwitch(boolean smartSwitch) {
        this.smartSwitch = smartSwitch;
    }
}